from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.ext.declarative import declarative_base



class Engine:
    def __init__(self, conn):
        self.engine = create_engine(conn, encoding='utf8')

    def excute(self, sql, **kwargs):
        """含commit操作，返回<class 'sqlalchemy.engine.result.ResultProxy'>"""
        return self.engine.execute(sql, **kwargs)

    def fetchall(self, sql):
        return self.execute(sql).fetchall()  # 能解释CLOB

    def fetchone(self, sql, n=999999):
        self.yellow(sql)
        result = self.execute(sql)
        for _ in range(n):
            one = result.fetchone()
            if one:
                yield one  # <class 'sqlalchemy.engine.result.RowProxy'>

    def fetchone_dt(self, sql, n=999999):
        self.yellow(sql)
        result = self.execute(sql)
        columns = result.keys()
        length = len(columns)
        for _ in range(n):
            one = result.fetchone()
            if one:
                yield {columns[i]: one[i] for i in range(length)}

    @staticmethod
    def yellow(x):
        print('\033[033m{}\033[0m'.format(x))

    @staticmethod
    def red(x):
        print('\033[031m{}\033[0m'.format(x))

class ORM(Engine):

    def __init__(self, conn, Base):
        super().__init__(conn)
        _Session = sessionmaker(bind=self.engine)
        self.session = _Session()
        self.Base = Base
        self.create_all()

    def __del__(self):
        self.session.close()

    def create_all(self):
        """创建所有表"""
        self.Base.metadata.create_all(bind=self.engine)

    def add(self, dt):
        """插入"""
        self.session.add(dt)  # 添加到ORM对象
        self.session.commit()  # 提交

    def query_all(self, Table):
        return self.session.query(Table).all()

    def update(self, Table, condition, dt):
        """有则更新，没则插入"""
        q = self.session.query(Table).filter_by(**condition)
        if q.all():
            q.update(dt)  # 更新
            self.session.commit()  # 提交
        else:
            self.add(Table, dt)

